package cn.iocoder.yudao.module.ai.websocket;

import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.framework.websocket.core.listener.WebSocketMessageListener;
import cn.iocoder.yudao.framework.websocket.core.message.JsonWebSocketMessage;
import cn.iocoder.yudao.module.ai.controller.admin.voice.vo.AiVoiceTranscriptionRespVO;
import cn.iocoder.yudao.module.ai.controller.admin.voice.vo.AiVoiceWebSocketMessage;
import cn.iocoder.yudao.module.ai.service.voice.AiVoiceService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import jakarta.annotation.Resource;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * AI 语音转文字 WebSocket 消息监听器
 *
 * @author 芋道源码
 */
@Component
@Slf4j
public class AiVoiceWebSocketMessageListener implements WebSocketMessageListener<AiVoiceWebSocketMessage> {

    @Resource
    private AiVoiceService aiVoiceService;

    /**
     * 消息类型
     */
    private static final String MESSAGE_TYPE = "ai-voice-transcription";

    /**
     * 存储每个会话的音频数据缓冲区
     * key: sessionId, value: 累积的音频数据
     */
    private final Map<String, StringBuilder> audioBufferMap = new ConcurrentHashMap<>();

    @Override
    public void onMessage(WebSocketSession session, AiVoiceWebSocketMessage message) {
        String sessionId = session.getId();
        log.info("[onMessage] 收到语音消息, sessionId: {}, type: {}", sessionId, message.getType());

        try {
            switch (message.getType()) {
                case "start":
                    handleStart(session);
                    break;
                case "audio":
                    handleAudio(session, message);
                    break;
                case "stop":
                    handleStop(session);
                    break;
                default:
                    log.warn("[onMessage] 未知的消息类型: {}", message.getType());
                    sendError(session, "未知的消息类型");
            }
        } catch (Exception e) {
            log.error("[onMessage] 处理语音消息失败", e);
            sendError(session, "处理失败: " + e.getMessage());
        }
    }

    /**
     * 处理开始识别消息
     */
    private void handleStart(WebSocketSession session) {
        String sessionId = session.getId();
        // 初始化音频缓冲区
        audioBufferMap.put(sessionId, new StringBuilder());
        
        log.info("[handleStart] 开始语音识别, sessionId: {}", sessionId);
        
        // 发送确认消息
        AiVoiceTranscriptionRespVO resp = AiVoiceTranscriptionRespVO.builder()
                .text("语音识别已开始...")
                .finished(false)
                .build();
        sendMessage(session, resp);
    }

    /**
     * 处理音频数据消息
     */
    private void handleAudio(WebSocketSession session, AiVoiceWebSocketMessage message) {
        String sessionId = session.getId();
        audioBufferMap.put(sessionId, new StringBuilder());
        StringBuilder buffer = audioBufferMap.get(sessionId);
        
        if (buffer == null) {
            log.warn("[handleAudio] 未找到音频缓冲区, sessionId: {}", sessionId);
            sendError(session, "请先发送 start 消息");
            return;
        }
        
        // 累积音频数据
        buffer.append(message.getData());
        
        log.debug("[handleAudio] 累积音频数据, sessionId: {}, 当前长度: {}", sessionId, buffer.length());
        
        // 可以选择实时转换，或者等待 stop 消息后一次性转换
        // 这里我们选择在收到 stop 消息后再转换，以保证音频的完整性
    }

    /**
     * 处理停止识别消息
     */
    private void handleStop(WebSocketSession session) {
        String sessionId = session.getId();
        StringBuilder buffer = audioBufferMap.remove(sessionId);
        
        if (buffer == null || buffer.isEmpty()) {
            log.warn("[handleStop] 未找到音频数据, sessionId: {}", sessionId);
            sendError(session, "没有可转换的音频数据");
            return;
        }
        
        log.info("[handleStop] 停止语音识别并开始转换, sessionId: {}, 音频数据长度: {}", sessionId, buffer.length());
        
        try {
            // 获取音频格式和语言
            String format = session.getAttributes().get("format") != null 
                    ? session.getAttributes().get("format").toString() : "webm";
            String language = session.getAttributes().get("language") != null 
                    ? session.getAttributes().get("language").toString() : "zh";
            
            // 调用语音转文字服务
            String audioData = buffer.toString();
            String text = aiVoiceService.transcribeAudioBase64(audioData, format, language);
            
            // 发送转换结果
            AiVoiceTranscriptionRespVO resp = AiVoiceTranscriptionRespVO.builder()
                    .text(text)
                    .finished(true)
                    .build();
            sendMessage(session, resp);
            
            log.info("[handleStop] 语音转换完成, sessionId: {}, 文本长度: {}", sessionId, text.length());
            
        } catch (Exception e) {
            log.error("[handleStop] 语音转文字失败", e);
            sendError(session, "语音转文字失败: " + e.getMessage());
        }
    }

    /**
     * 发送响应消息
     */
    private void sendMessage(WebSocketSession session, AiVoiceTranscriptionRespVO resp) {
        if (session == null || !session.isOpen()) {
            log.warn("[sendMessage] Session 为空或已关闭");
            return;
        }
        
        try {
            // 构造 WebSocket 消息
            JsonWebSocketMessage message = new JsonWebSocketMessage();
            message.setType(MESSAGE_TYPE);
            message.setContent(JsonUtils.toJsonString(resp));
            
            // 发送消息
            String payload = JsonUtils.toJsonString(message);
            session.sendMessage(new TextMessage(payload));
            
            log.debug("[sendMessage] 发送消息成功, sessionId: {}", session.getId());
        } catch (IOException e) {
            log.error("[sendMessage] 发送消息失败", e);
        }
    }

    /**
     * 发送错误消息
     */
    private void sendError(WebSocketSession session, String error) {
        AiVoiceTranscriptionRespVO resp = AiVoiceTranscriptionRespVO.builder()
                .error(error)
                .finished(true)
                .build();
        sendMessage(session, resp);
    }

    @Override
    public String getType() {
        return MESSAGE_TYPE;
    }

}

